package org.happy.tester.agent;

import com.alibaba.fastjson.JSONObject;
import javassist.*;
import org.happy.tester.common.http.req.PointcutRequest;
import org.happy.tester.common.model.PointcutDO;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

import static org.happy.tester.common.constant.CommonConstant.POINTCUT_JOINER;

public class TargetClassTransFormer implements ClassFileTransformer {

    private PointcutRequest pointcutRequest;

    private Map<String, List<PointcutDO>> pointcutDOMap;

    TargetClassTransFormer(PointcutRequest pointcutRequest, Map<String, List<PointcutDO>> pointcutDOMap) {
        this.pointcutRequest = pointcutRequest;
        this.pointcutDOMap = pointcutDOMap;
    }

    private static final String[] NECESSARY_PACKAGES = new String[]{"org.happy.tester.common.http", "com.alibaba.fastjson.JSONObject"};

    /**
     * instrumentation 有局限性
     **/
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        className = className.replace("/", ".");
        if (pointcutDOMap.containsKey(className)) {
            List<PointcutDO> pointcutDOList = pointcutDOMap.get(className);
            ClassPool pool = ClassPool.getDefault();
            try {
                pool.insertClassPath(classBeingRedefined.getResource("/").getPath());
                System.out.println(JSONObject.toJSONString(classBeingRedefined.getResource("/").getPath()));
                for (String necessaryPackage : NECESSARY_PACKAGES) {
                    pool.importPackage(necessaryPackage);
                }
                CtClass ctClass = pool.get(className);
                if (ctClass.isFrozen()) {
                    ctClass.defrost();
                }
                for (PointcutDO pointcutDO : pointcutDOList) {
                    CtClass[] paramCtClasses = pool.get(pointcutDO.getParamsTypeArray());
                    CtMethod method = ctClass.getDeclaredMethod(pointcutDO.getMethodName(), paramCtClasses);
                    String resultClassName = method.getReturnType().getClassFile().getName();
                    String pointcutMiddle = concatPointcutMiddle(pointcutDO);
                    String beforeMethod = getBeforeMethod(pointcutMiddle, resultClassName);
                    String afterMethod = getAfterMethod(pointcutMiddle);
                    method.insertBefore(beforeMethod);
                    method.insertAfter(afterMethod, true);
                }
                ctClass.detach();
                return ctClass.toBytecode();
            } catch (NotFoundException | CannotCompileException | IOException e) {
                e.printStackTrace();
            }
        }
        return classfileBuffer;
    }

    private String getBeforeMethod(String pointcutMiddle, String resultClassName) {
        StringBuffer beforeBuffer = new StringBuffer();
        beforeBuffer.append("\nCatcherClient client = new CatcherClient(\"" + pointcutMiddle + "\");\n")
                .append("boolean needMock = client.needMock();\n")
                .append("if (needMock) {\n")
                .append("Object result = client.getMockResult(\"" + resultClassName + "\");\n")
                .append("return ($r)result;\n")
                .append("}\n");
        return beforeBuffer.toString();
    }

    private String getAfterMethod(String pointcutMiddle) {
        StringBuffer afterBuffer = new StringBuffer();
        afterBuffer.append("\nCatcherClient client = new CatcherClient(\"" + pointcutMiddle + "\");\n")
                .append("boolean needCollect = client.needCollect();\n")
                .append("if (needCollect) {\n")
                .append("client.pushCollectResult($_);\n")
                .append("}\n");
        return afterBuffer.toString();
    }

    private String concatPointcutMiddle(PointcutDO pointcutDO) {
        String paramTypes = null;
        if (pointcutDO.getParamsTypeArray() != null) {
            paramTypes = Arrays.toString(pointcutDO.getParamsTypeArray());
            paramTypes = paramTypes.substring(1, paramTypes.length() - 1);
        }
        String pointcutMiddle = POINTCUT_JOINER.join(pointcutDO.getClassName(),
                pointcutDO.getMethodName(),
                paramTypes);
        return pointcutMiddle;
    }
}